from gensim.models.word2vec import Word2Vec
from gensim.models.word2vec import LineSentence
import logging

#训练模型
logging.basicConfig(format='%(asctime)s : %(levelname)s : %(message)s')
logging.root.setLevel(level=logging.INFO)

def gensim_model(segFile, vecFile, sg = 1, dim=64):
    # segFile = "D:/data/thudata_seg.txt"
    # vecFile = "gensim64_skipgram_thudata.vec"
    #
    sentence = LineSentence(segFile)
    logging.info("sentence finished")

    #use cbow sg = 0 or skipgram sg = 1
    model = Word2Vec(sentence, size=dim, min_count=1, sg=sg)
    logging.info("model finished")
    model.wv.save_word2vec_format('model/' + vecFile)


gensim_model("D:/data/videodata_seg.txt", "gensim64_skipgram_webvideo.vec", 1, 64)
gensim_model("D:/data/thudata_seg.txt", "gensim64_skipgram_thudata.vec", 1, 64)
gensim_model("D:/data/merged_seg.txt", "gensim64_skipgram_merge.vec", 1, 64)



# 载入模型
from gensim.models.keyedvectors import KeyedVectors
from gensim.models.fasttext import FastText

model = KeyedVectors.load_word2vec_format("model/gensim64_skipgram_merge" + ".vec")
model.most_similar("中共")

model = FastText.load_fasttext_format("model/fastetxt64_skipgram_merge" + ".bin")
model.most_similar("中共")


# cmd 训练fasttext
# fasttext cbow -input webvideodata20171020_90.utf8_fixed_seg.utf8 -output fasttext64_cbow -minCount 1 -dim 64


